//
// Copyright (c) 2007 Rodrigo B. de Oliveira & Jb Evain
//
// Permission is hereby granted, free of charge, to any person obtaining
// a copy of this software and associated documentation files (the
// "Software"), to deal in the Software without restriction, including
// without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to
// the following conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
//

using System;
using System.Collections;
using System.IO;
using Milo.Compilation;
using Milo.Ast;
using Milo.Compilation.Steps;
using Milo.Compilation.TypeSystem;
using Mono.Cecil.Cil;
using Cecil = Mono.Cecil;

namespace Milo.Compilation
{
	public class EmitAssembly : AbstractVisitorCompilerStep
	{
		private static readonly object DefinitionKey = new object();

		private Cecil.AssemblyDefinition _assembly;
		private CilWorker _cil;
		private Stack _stack = new Stack();

		public override void Run()
		{
			if (Errors.Count > 0) return;

			SetUpAssemblyDefinition();
			DefineTypes();
			DefineMembers();
			EmitTypes();
			DefineEntryPoint();
			SaveAssemblyDefinition();
		}

		private void DefineMembers()
		{
			CompileUnit.Accept(new DefineMembersDispatcher(this));
		}

		private void DefineEntryPoint()
		{
			if (CompilerOutputType.Library == Parameters.OutputType) return;

			Method method = ContextAnnotations.GetEntryPoint(Context);
			if (null != method)
			{
				_assembly.EntryPoint = GetMethodDefinition(method);
			}
			else
			{
				Errors.Add(CompilerErrorFactory.NoEntryPoint());
			}
		}

		private void SaveAssemblyDefinition()
		{
			if (Context.Parameters.GenerateInMemory)
			{
				byte[] assembly;
				Cecil.AssemblyFactory.SaveAssembly(_assembly, out assembly);
				Context.GeneratedAssembly = System.Reflection.Assembly.Load(assembly);
			}
			else
			{
				Cecil.AssemblyFactory.SaveAssembly(_assembly, GetOutputAssemblyFileName());
			}
		}

		private string GetOutputAssemblyFileName()
		{
			if (Parameters.OutputAssembly != string.Empty) return Parameters.OutputAssembly;
			return BuildOutputAssemblyName(GetAssemblyName());
		}

		string BuildOutputAssemblyName(string fname)
		{
			if (!Path.HasExtension(fname))
			{
				if (CompilerOutputType.Library == Parameters.OutputType)
				{
					fname += ".dll";
				}
				else
				{
					fname += ".exe";
				}
			}
			return Path.GetFullPath(fname);
		}

		private Cecil.AssemblyKind GetAssemblyKind()
		{
			switch (Parameters.OutputType)
			{
				case CompilerOutputType.ConsoleApplication:
					return Cecil.AssemblyKind.Console;
				case CompilerOutputType.Library:
					return Cecil.AssemblyKind.Dll;
				case CompilerOutputType.WindowsApplication:
					return Cecil.AssemblyKind.Windows;
			}
			throw new ArgumentException();
		}

		private string GetAssemblyName()
		{
			string fname = Parameters.OutputAssembly;
			if (0 == fname.Length)
			{
				return CompileUnit.Modules[0].Name;
			}
			return Path.GetFileNameWithoutExtension(fname);
		}

		private void EmitTypes()
		{
			CompileUnit.Accept(this);
		}

		private Cecil.TypeDefinition EmitTypeDefinition(TypeDefinition node)
		{
			IType type = (IType)GetEntity(node);
			Cecil.TypeDefinition typeDef = GetTypeDefinition(node);
			foreach (IType iface in type.GetInterfaces())
			{
				typeDef.Interfaces.Add(GetTypeReference(iface));
			}

			Visit(node.Members);
			return typeDef;
		}

		public override void OnClassDefinition(ClassDefinition node)
		{
			IType type = (IType)GetEntity(node);
			Cecil.TypeDefinition typeDef = EmitTypeDefinition(node);
			typeDef.BaseType = GetTypeReference(type.BaseType);
		}

		public override void OnInterfaceDefinition(InterfaceDefinition node)
		{
			EmitTypeDefinition(node);
		}

		public override void OnConstructor(Constructor node)
		{
			GetTypeDefinition(node.DeclaringType).Constructors.Add(EmitMethod(node));
		}

		public override void OnMethod(Method node)
		{
			GetTypeDefinition(node.DeclaringType).Methods.Add(EmitMethod(node));
		}

		public override void OnField(Field node)
		{
			GetTypeDefinition(node.DeclaringType).Fields.Add(GetFieldDefinition(node));
		}

		private Cecil.FieldDefinition DefineField(Field node)
		{
			IField field = (IField)GetEntity(node);
			Cecil.FieldDefinition fieldDef =
				new Cecil.FieldDefinition(node.Name, GetTypeReference(field.Type), GetFieldAttributes(node));

			SetDefinition(node, fieldDef);

			return fieldDef;
		}

		private static Cecil.FieldAttributes GetFieldAttributes(Field node)
		{
			Cecil.FieldAttributes attributes = GetAccessibilityAttribute(node);
			if (node.IsStatic) attributes |= Cecil.FieldAttributes.Static;
			if (node.IsFinal) attributes |= Cecil.FieldAttributes.InitOnly;
			return attributes;
		}

		private Cecil.MethodDefinition EmitMethod(Method node)
		{
			IMethod method = GetEntity(node);
			Cecil.MethodDefinition methodDef =
				new Cecil.MethodDefinition(GetMethodName(node), GetMethodAttributes(node), GetTypeReference(method.ReturnType));

			SetDefinition(node, methodDef);

			DefineParameters(node, methodDef);

			if (methodDef.HasBody)
			{
				DefineLocalVariables(node, methodDef);
				_cil = methodDef.Body.CilWorker;

				Visit(node.Body);

				_cil.Emit(OpCodes.Ret);
			}

			return methodDef;
		}

		private void DefineLocalVariables(Method node, Cecil.MethodDefinition method)
		{
			if (node.Locals.Count == 0) return;

			method.Body.InitLocals = true;

			for (int i = 0; i < node.Locals.Count; i++)
			{
				method.Body.Variables.Add(DefineLocalVariable(node.Locals[i], i, method));
			}
		}

		private void DefineParameters(Method node, Cecil.MethodDefinition method)
		{
			foreach (ParameterDeclaration declaration in node.Parameters)
			{
				method.Parameters.Add(DefineParameter(declaration));
			}
		}

		private VariableDefinition DefineLocalVariable(Local local, int index, Cecil.MethodDefinition method)
		{
			InternalLocal entity = (InternalLocal) GetEntity(local);
			VariableDefinition variable = new VariableDefinition(local.Name, index, method, GetTypeReference(entity.Type));
			SetDefinition(local, variable);
			return variable;
		}

		private Cecil.ParameterDefinition DefineParameter(ParameterDeclaration declaration)
		{
			InternalParameter entity = (InternalParameter) GetEntity(declaration);
			Cecil.ParameterDefinition parameter =
				new Cecil.ParameterDefinition(declaration.Name, entity.Index, Cecil.ParameterAttributes.None,
				                              GetTypeReference(entity.Type));
			SetDefinition(declaration, parameter);
			return parameter;
		}

		private static string GetMethodName(Method node)
		{
			if (IsConstructor(node))
			{
				if (node.IsStatic) return Cecil.MethodDefinition.Cctor;
				return Cecil.MethodDefinition.Ctor;
			}
			return node.Name;
		}

		public override void LeaveExpressionStatement(ExpressionStatement node)
		{
			EmitPopIfNeeded(PopType());
			AssertStackIsEmpty();
		}

		private void AssertStackIsEmpty()
		{
			if (_stack.Count == 0) return;
			throw new InvalidOperationException("stack should be empty");
		}

		public override void OnMethodInvocationExpression(MethodInvocationExpression node)
		{
			if (IsBuiltinFunction(node))
			{
				EmitBuiltinInvocation(node);
				return;
			}

			if (IsSuperInvocation(node))
			{
				EmitSuperInvocation(node);
			}
			else
			{
				EmitInvocation(node);
			}
		}

		private void EmitBuiltinInvocation(MethodInvocationExpression node)
		{
			BuiltinFunction builtin = (BuiltinFunction) GetEntity(node.Target);
			switch (builtin.FunctionType)
			{
				case BuiltinFunctionType.Eval:
					EmitEval(node);
					return;
			}
			UnsupportedExpression(node);
		}

		private void EmitEval(MethodInvocationExpression node)
		{
			int allButLast = node.Arguments.Count-1;
			for (int i = 0; i < allButLast; ++i)
			{
				Visit(node.Arguments[i]);
				EmitPopIfNeeded(PopType());
			}
			Visit(node.Arguments[-1]);
		}

		private static bool IsBuiltinFunction(MethodInvocationExpression node)
		{
			return TypeSystemHelpers.IsBuiltinFunction(node);
		}

		private void EmitPopIfNeeded(IType typeOnStack)
		{
			if (TypeSystemServices.VoidType == typeOnStack) return;

			_cil.Emit(OpCodes.Pop);
		}

		private static bool IsSuperInvocation(MethodInvocationExpression node)
		{
			return node.Target.NodeType == NodeType.SuperLiteralExpression;
		}

		private void EmitSuperInvocation(MethodInvocationExpression node)
		{
			EmitSelf();
			EmitArguments(node);

			IMethod method = (IMethod)GetEntity(node.Target);
			_cil.Emit(OpCodes.Call, GetMethodReference(method));

			PushType(method.ReturnType);
		}

		private void EmitArguments(MethodInvocationExpression node)
		{
			foreach (Expression arg in node.Arguments)
			{
				Visit(arg);
				PopType();
			}
		}

		private void EmitInvocation(MethodInvocationExpression node)
		{
			IEntity entity = GetEntity(node.Target);
			if (EntityType.Constructor == entity.EntityType)
			{
				EmitConstructorInvocation(node, (IConstructor)entity);
			}
			else
			{
				EmitMethodInvocation(node, (IMethod) entity);
			}
		}

		private void EmitConstructorInvocation(MethodInvocationExpression node, IConstructor constructor)
		{
			EmitArguments(node);
			_cil.Emit(OpCodes.Newobj, GetMethodReference(constructor));
			PushType(constructor.DeclaringType);
		}

		private void EmitMethodInvocation(MethodInvocationExpression node, IMethod method)
		{
			if (!method.IsStatic)
			{
				EmitTargetInstance(node.Target);
			}

			EmitArguments(node);

			_cil.Emit(method.IsStatic ? OpCodes.Call : OpCodes.Callvirt, GetMethodReference(method));

			PushType(method.ReturnType);
		}

		private static Expression GetTargetInstance(Expression expression)
		{
			return ((MemberReferenceExpression)expression).Target;
		}

		public override void OnSelfLiteralExpression(SelfLiteralExpression node)
		{
			EmitSelf();
			PushExpressionType(node);
		}

		private void EmitSelf()
		{
			_cil.Emit(OpCodes.Ldarg_0);
		}

		public override void LeaveReturnStatement(ReturnStatement node)
		{
			if (null == node.Expression) return;
			PopType();
		}

		public override void OnCastExpression(CastExpression node)
		{
			Visit(node.Target);

			IType actualType = PopType();
			IType expectedType = (IType)GetEntity(node.Type);
			if (TypeSystemServices.RequiresBoxing(expectedType, actualType))
			{
				_cil.Emit(OpCodes.Box, GetTypeReference(actualType));
			}
			else
			{
				_cil.Emit(OpCodes.Castclass, GetTypeReference(node.Type));
			}

			PushType(expectedType);
		}

		private Cecil.TypeReference GetTypeReference(TypeReference typeReference)
		{
			return GetTypeReference((IType)GetEntity(typeReference));
		}

		public override void OnNullLiteralExpression(NullLiteralExpression node)
		{
			_cil.Emit(OpCodes.Ldnull);
			PushType(null);
		}

		public override void OnBinaryExpression(BinaryExpression node)
		{
			switch (node.Operator)
			{
				case BinaryOperatorType.Assign:
					EmitAssignment(node);
					PushVoid();
					return;
			}
			UnsupportedExpression(node);
		}

		private void EmitAssignment(BinaryExpression node)
		{
			IEntity lvalue = GetEntity(node.Left);
			switch (lvalue.EntityType)
			{
				case EntityType.Field:
					EmitFieldAssignment(node);
					return;
				case EntityType.Local:
					EmitLocalAssignment(node);
					return;
			}
			UnsupportedExpression(node);
		}

		private void EmitLocalAssignment(BinaryExpression node)
		{
			Visit(node.Right);
			PopType();
			EmitStoreLocal(node);
		}

		private void EmitStoreLocal(BinaryExpression node)
		{
			VariableDefinition variable = GetVariableDefinition(((InternalLocal)GetEntity(node.Left)).Local);
			_cil.Emit(OpCodes.Stloc, variable);
		}

		private void EmitFieldAssignment(BinaryExpression node)
		{
			IField field = (IField)GetEntity(node.Left);
			if (!field.IsStatic)
			{
				EmitTargetInstance(node.Left);
			}
			Visit(node.Right);
			PopType();
			EmitStoreField(field);
		}

		private void EmitStoreField(IField field)
		{
			Cecil.FieldReference reference = GetFieldReference(field);
			if (field.IsStatic)
			{
				_cil.Emit(OpCodes.Stsfld, reference);
			}
			else
			{
				_cil.Emit(OpCodes.Stfld, reference);
			}
		}

		private void EmitTargetInstance(Expression expression)
		{
			Visit(GetTargetInstance(expression));
			PopType();
		}

		private static void UnsupportedExpression(Expression node)
		{
			throw new Exception("Unsupported expression: " + node);
		}

		public override void OnIntegerLiteralExpression(IntegerLiteralExpression node)
		{
			_cil.Emit(OpCodes.Ldc_I4, (int)node.Value);
			PushExpressionType(node);
		}

		public override void OnStringLiteralExpression(StringLiteralExpression node)
		{
			_cil.Emit(OpCodes.Ldstr, node.Value);
			PushExpressionType(node);
		}

		public override void OnReferenceExpression(ReferenceExpression node)
		{
			IEntity entity = GetEntity(node);
			switch (entity.EntityType)
			{
				case EntityType.Parameter:
					EmitLoadParameter(node);
					break;
				case EntityType.Local:
					EmitLoadVariable(node);
					break;
				default:
					UnsupportedExpression(node);
					return;
			}
			PushExpressionType(node);
		}

		private void EmitLoadVariable(ReferenceExpression node)
		{
			VariableDefinition variable = GetVariableDefinition(((InternalLocal)GetEntity(node)).Local);
			switch (variable.Index)
			{
				case 0:
					_cil.Emit(OpCodes.Ldloc_0);
					return;
				case 1:
					_cil.Emit(OpCodes.Ldloc_1);
					return;
			}
		}

		private void EmitLoadParameter(ReferenceExpression node)
		{
			InternalParameter parameter = (InternalParameter)GetEntity(node);
			switch (parameter.Index)
			{
				case 0:
					_cil.Emit(OpCodes.Ldarg_0);
					return;
				case 1:
					_cil.Emit(OpCodes.Ldarg_1);
					return;
			}
			UnsupportedExpression(node);
		}

		public override void OnMemberReferenceExpression(MemberReferenceExpression node)
		{
			IField field = (IField) GetEntity(node);
			if (field.IsStatic)
			{
				_cil.Emit(OpCodes.Ldsfld, GetFieldReference(field));
			}
			else
			{
				EmitTargetInstance(node);
				_cil.Emit(OpCodes.Ldfld, GetFieldReference(field));
			}
			PushExpressionType(node);
		}

		private void PushExpressionType(Expression node)
		{
			PushType(GetExpressionType(node));
		}

		private static Cecil.TypeAttributes GetTypeAttributes(TypeDefinition type)
		{
			Cecil.TypeAttributes attributes = GetAccessibilityAttribute(type);
			if (type.IsAbstract) attributes |= Cecil.TypeAttributes.Abstract;
			return attributes;
		}

		private static Cecil.MethodAttributes GetMethodAttributes(Method node)
		{
			Cecil.MethodAttributes attributes = GetAccessibilityAttribute(node);
			if (node.IsStatic)
			{
				attributes |= Cecil.MethodAttributes.Static;
			}
			else
			{
				if (node.IsVirtual || node.IsOverride)
				{
					attributes |= Cecil.MethodAttributes.Virtual;
					if (!node.IsOverride)
					{
						attributes |= Cecil.MethodAttributes.NewSlot;
					}
				}
				if (node.IsAbstract)
				{
					attributes |= Cecil.MethodAttributes.Abstract;
					attributes |= Cecil.MethodAttributes.Virtual;
				}
			}
			if (IsConstructor(node))
			{
				attributes |= Cecil.MethodAttributes.SpecialName;
				attributes |= Cecil.MethodAttributes.RTSpecialName;
			}
			return attributes;
		}

		private static bool IsConstructor(Method node)
		{
			return node.NodeType == NodeType.Constructor;
		}

		private static Cecil.TypeAttributes GetAccessibilityAttribute(TypeDefinition node)
		{
			if (node.IsPublic) return Cecil.TypeAttributes.Public;
			return Cecil.TypeAttributes.NotPublic;
		}

		private static Cecil.FieldAttributes GetAccessibilityAttribute(Field node)
		{
			if (node.IsPublic) return Cecil.FieldAttributes.Public;
			if (node.IsInternal) return Cecil.FieldAttributes.Assembly;
			if (node.IsProtected) return Cecil.FieldAttributes.Family;
			return Cecil.FieldAttributes.Private;
		}

		private static Cecil.MethodAttributes GetAccessibilityAttribute(Method node)
		{
			if (node.IsPublic) return Cecil.MethodAttributes.Public;
			if (node.IsInternal) return Cecil.MethodAttributes.Assem;
			if (node.IsProtected) return Cecil.MethodAttributes.Family;
			return Cecil.MethodAttributes.Private;
		}

		private Cecil.TypeReference GetTypeReference(IType type)
		{
			if (type.IsArray) return GetArrayTypeReference((IArrayType) type);
			ExternalType external = type as ExternalType;
			if (null != external) return Import(external.ActualType);
			return GetTypeDefinition(((AbstractInternalType)type).TypeDefinition);
		}

		private Cecil.TypeReference GetArrayTypeReference(IArrayType type)
		{
			return new Cecil.ArrayType(GetTypeReference(type.GetElementType()));
		}

		private Cecil.MethodReference GetMethodReference(IMethodBase method)
		{
			ExternalMethod external = method as ExternalMethod;
			if (null != external) return Import(external.MethodInfo);
			return GetMethodDefinition(((InternalMethod)method).Method);
		}

		private Cecil.FieldReference GetFieldReference(IField field)
		{
			ExternalField external = field as ExternalField;
			if (null != external) return Import(external.FieldInfo);
			return GetFieldDefinition(((InternalField) field).Field);
		}

		private Cecil.FieldReference Import(System.Reflection.FieldInfo field)
		{
			return _assembly.MainModule.Import(field);
		}

		private Cecil.MethodReference Import(System.Reflection.MethodBase method)
		{
			return _assembly.MainModule.Import(method);
		}

		private Cecil.TypeReference Import(System.Type type)
		{
			return _assembly.MainModule.Import(type);
		}

		class DefineTypesDispatcher : DepthFirstVisitor
		{
			private EmitAssembly _emitter;

			public DefineTypesDispatcher(EmitAssembly emitter)
			{
				_emitter = emitter;
			}

			public override void OnClassDefinition(ClassDefinition node)
			{
				_emitter.DefineClass(node);
			}

			public override void OnInterfaceDefinition(InterfaceDefinition node)
			{
				_emitter.DefineInterface(node);
			}
		}

		class DefineMembersDispatcher : DepthFirstVisitor
		{
			private EmitAssembly _emitter;

			public DefineMembersDispatcher(EmitAssembly emitter)
			{
				_emitter = emitter;
			}

			public override void OnField(Field node)
			{
				_emitter.DefineField(node);
			}
		}

		private void DefineTypes()
		{
			CompileUnit.Accept(new DefineTypesDispatcher(this));
		}

		private void SetUpAssemblyDefinition()
		{
			_assembly = Cecil.AssemblyFactory.DefineAssembly(GetAssemblyName(), GetAssemblyKind());
		}

		public void DefineClass(ClassDefinition node)
		{
			DefineTypeDefinition(node);
		}

		public void DefineInterface(InterfaceDefinition node)
		{
			Cecil.TypeDefinition iface = DefineTypeDefinition(node);
			iface.Attributes |= Cecil.TypeAttributes.Interface;
			iface.Attributes |= Cecil.TypeAttributes.Abstract;
		}

		private Cecil.TypeDefinition DefineTypeDefinition(TypeDefinition node)
		{
			Cecil.TypeDefinition type = new Cecil.TypeDefinition(node.Name, GetNamespace(node), GetTypeAttributes(node), null);
			SetDefinition(node, type);
			_assembly.MainModule.Types.Add(type);
			return type;
		}

		private static void SetDefinition(Node node, object definition)
		{
			node.Annotate(DefinitionKey, definition);
		}

		private static Cecil.FieldDefinition GetFieldDefinition(Field node)
		{
			return (Cecil.FieldDefinition)GetDefinition(node);
		}

		private static object GetDefinition(Node node)
		{
			object definition = node[DefinitionKey];
			if (null == definition) throw CompilerErrorFactory.InternalError(node, "Unbound node", null);
			return definition;
		}

		private static Cecil.MethodDefinition GetMethodDefinition(Method node)
		{
			return (Cecil.MethodDefinition)GetDefinition(node);
		}

		private static Cecil.TypeDefinition GetTypeDefinition(TypeDefinition node)
		{
			return (Cecil.TypeDefinition)GetDefinition(node);
		}

		private static VariableDefinition GetVariableDefinition(Local node)
		{
			return (VariableDefinition)GetDefinition(node);
		}

		private static string GetNamespace(TypeDefinition node)
		{
			NamespaceDeclaration declaration = node.EnclosingNamespace;
			if (null != declaration) return declaration.Name;
			return string.Empty;
		}

		private void PushVoid()
		{
			PushType(TypeSystemServices.VoidType);
		}

		private void PushType(IType type)
		{
			_stack.Push(type);
		}

		private IType PopType()
		{
			return (IType)_stack.Pop();
		}
	}
}